package com.srtp.secureprotocol.cipherstream;

import com.srtp.secureprotocol.util.Util;

import java.io.IOException;
import java.io.OutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.Mac;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * 加密文件输出流，继承自摘要输出流
 * 提供文件加密功能，支持AES/CBC/PKCS5Padding加密模式
 * 同时计算HMAC-SHA256消息认证码确保数据完整性
 */
public class CipherOutputStream extends DigestingOutputStream {

    private final Cipher cipher; // AES加密器
    private final Mac mac;    // HMAC-SHA256消息认证码计算器

    /**
     * 加密输出流构造函数
     *
     * @param combinedKeyMaterial 组合密钥材料（加密密钥+MAC密钥）
     * @param iv                  初始化向量（IV），如果为null则自动生成
     * @param outputStream        底层输出流
     * @throws IOException 流操作错误
     */
    public CipherOutputStream(byte[] combinedKeyMaterial,
                              byte[] iv,
                              OutputStream outputStream)
            throws IOException {
        super(outputStream);
        try {
            // 初始化加密器和MAC计算器
            this.cipher = initializeCipher();
            this.mac = initializeMac();

            // 分割组合密钥为加密密钥和MAC密钥
            byte[][] keyParts = Util.split(combinedKeyMaterial, 32, 32);

            // 配置加密器（自动生成IV或使用提供的IV）
            if (iv == null) {
                this.cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(keyParts[0], "AES"));
            } else {
                this.cipher.init(Cipher.ENCRYPT_MODE,
                        new SecretKeySpec(keyParts[0], "AES"),
                        new IvParameterSpec(iv));
            }

            // 初始化MAC计算器
            this.mac.init(new SecretKeySpec(keyParts[1], "HmacSHA256"));

            // 将IV加入MAC计算并写入输出流
            mac.update(cipher.getIV());
            super.write(cipher.getIV());
        } catch (InvalidKeyException | InvalidAlgorithmParameterException e) {
            throw new AssertionError("加密初始化失败", e);
        }
    }

    @Override
    public void write(byte[] buffer) throws IOException {
        write(buffer, 0, buffer.length);
    }

    @Override
    public void write(byte[] buffer, int offset, int length) throws IOException {
        // 加密数据块
        byte[] ciphertext = cipher.update(buffer, offset, length);

        if (ciphertext != null) {
            // 更新MAC并写入加密数据
            mac.update(ciphertext);
            super.write(ciphertext);
        }
    }

    @Override
    public void write(int b) {
        // 未实现单字节写入（因加密效率考虑）
        throw new AssertionError("单字节写入未实现");
    }

    /**
     * 刷新缓冲区并完成加密
     * <p>
     * 执行关键操作：
     * 1. 完成加密（处理最终块和填充）
     * 2. 计算最终MAC值
     * 3. 写入加密尾部和MAC
     * 4. 刷新底层流
     */
    @Override
    public void flush() throws IOException {
        try {
            // 完成加密（处理PKCS5填充）
            byte[] ciphertext = cipher.doFinal();
            // 计算完整数据的MAC
            byte[] auth = mac.doFinal(ciphertext);

            // 写入最终加密数据和MAC
            super.write(ciphertext);
            super.write(auth);

            // 刷新底层输出流
            super.flush();
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            throw new AssertionError("加密完成失败", e);
        }
    }

    /**
     * 初始化HMAC-SHA256计算器
     *
     * @return 配置好的MAC实例
     */
    private Mac initializeMac() {
        try {
            return Mac.getInstance("HmacSHA256");
        } catch (NoSuchAlgorithmException e) {
            throw new AssertionError("HMAC-SHA256算法不可用", e);
        }
    }

    /**
     * 初始化AES加密器
     *
     * @return 配置好的Cipher实例
     */
    private Cipher initializeCipher() {
        try {
            // 使用AES/CBC模式，PKCS5Padding填充方案
            return Cipher.getInstance("AES/CBC/PKCS5Padding");
        } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
            throw new AssertionError("AES加密器初始化失败", e);
        }
    }
}